package com.tfswx.ocrengine.util;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.tfswx.ocrengine.dto.*;
import com.tfswx.ocrengine.exception.AppException;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.pdfbox.io.MemoryUsageSetting;
import org.apache.pdfbox.multipdf.PDFMergerUtility;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * 对外提供接口识别结果封装工具
 * @author Ricky
 * @date 2021/1/11
 */
public class SbjgUtils {
    private SbjgUtils() {
    }

    /**
     * 封装响应结果数据
     *
     * @param ocrOutputDtoList
     * @param sfwjnr
     * @return
     */
    public static WjOcrSbjgGetOutputDTO packageResponseStr(
            String wjbh,
            String wjhz,
            List<OcrOutputDTO> ocrOutputDtoList,
            boolean sfwjwb,
            boolean sfwjxx,
            boolean sfwjnr) {
        // 封装参数文件识信息
        List<WjsbxxBean> wjsbxxBeanList = new ArrayList<>();
        StringBuilder pdfnrBuilder = null;
        if (sfwjwb) {
            pdfnrBuilder = new StringBuilder();
        }
        for (OcrOutputDTO ocrResult : ocrOutputDtoList) {
            WjsbxxBean wjsbxxBean = new WjsbxxBean();

            wjsbxxBeanList.add(wjsbxxBean);
            wjsbxxBean.setWjnr(ocrResult.getPdfStr());
            if (sfwjwb) {
                wjsbxxBean.setWjwb(ocrResult.getContent());
            }
            wjsbxxBean.setWjgd(ocrResult.getHeight());
            wjsbxxBean.setWjkd(ocrResult.getWidth());
            wjsbxxBean.setWjysgd(ocrResult.getOrgHeight());
            wjsbxxBean.setWjyskd(ocrResult.getOrgWidth());

            wjsbxxBean.setWjqxjd(ocrResult.getAngle());

            // 封装单页PDF识别后的文件识别结果信息
            if (sfwjxx) {
                WjjgxxBean wjjgxxBean = packageWjjgxxInfo(pdfnrBuilder, ocrResult);
                wjsbxxBean.setWjjgxx(wjjgxxBean);
            }
        }

        // 最终响应结果
        WjOcrSbjgGetOutputDTO wjOcrSbjgGetOutputDTO = new WjOcrSbjgGetOutputDTO();
        wjOcrSbjgGetOutputDTO.setWjsbxx(wjsbxxBeanList);
        wjOcrSbjgGetOutputDTO.setWjbh(wjbh);
        wjOcrSbjgGetOutputDTO.setWjhz(wjhz);
        wjOcrSbjgGetOutputDTO.setWjys(ocrOutputDtoList.size());

        if (sfwjnr) {
            try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
                PDFMergerUtility mergePdf = new PDFMergerUtility();
                ocrOutputDtoList.forEach(ocrOutputDTO -> mergePdf.addSource(new ByteArrayInputStream(Base64Util.getByteArray(ocrOutputDTO.getPdfStr()))));
                mergePdf.setDestinationStream(outputStream);
                mergePdf.mergeDocuments(MemoryUsageSetting.setupMainMemoryOnly());
                byte[] bytes = outputStream.toByteArray();
                wjOcrSbjgGetOutputDTO.setPdfnr(Base64Util.getBase64String(bytes));
                wjOcrSbjgGetOutputDTO.setWjmd5(DigestUtils.md5Hex(bytes));
            } catch (IOException e) {
                throw new AppException("合并pdf失败：" + e.getMessage(), e);
            }
        }

        if (pdfnrBuilder != null) {
            wjOcrSbjgGetOutputDTO.setTxtnr(pdfnrBuilder.toString());
        }

        return wjOcrSbjgGetOutputDTO;
    }

    /**
     * @param pdfnrBuilder
     * @param ocrResult
     * @return
     */
    private static WjjgxxBean packageWjjgxxInfo(StringBuilder pdfnrBuilder, OcrOutputDTO ocrResult) {

        List<WbxxBean> wbxxBeans = packageWbxxInfo(pdfnrBuilder, ocrResult);

        WjjgxxBean wjjgxxBean = new WjjgxxBean();
        wjjgxxBean.setWbsl(wbxxBeans.size());
        wjjgxxBean.setWbxx(wbxxBeans);

        wjjgxxBean.setSxsl(0);
        wjjgxxBean.setSxxx(new ArrayList<>(0));

        List<BgxxBean> bgxxBeans = packageBgxxInfo(ocrResult);

        wjjgxxBean.setBgsl(bgxxBeans.size());
        wjjgxxBean.setBgxx(bgxxBeans);

        return wjjgxxBean;
    }

    private static List<BgxxBean> packageBgxxInfo(OcrOutputDTO ocrResult) {
        List<BgxxBean> bgxxBeans = new LinkedList<>();

        JSONArray prismTablesInfo = ocrResult.getPrismTablesInfo();
        if (prismTablesInfo != null && !prismTablesInfo.isEmpty()) {
            for (int i = 0; i < prismTablesInfo.size(); i++) {
                JSONObject prismObj = prismTablesInfo.getJSONObject(i);
                JSONArray cellInfos = prismObj.getJSONArray("cellInfos");
                BgxxBean bgxxBean = new BgxxBean();
                List<DygxxBean> dygxxBeans = new ArrayList<>();
                for (int i1 = 0; i1 < cellInfos.size(); i1++) {
                    JSONObject cellInfo = cellInfos.getJSONObject(i1);
                    DygxxBean dygxxBean = new DygxxBean();
                    dygxxBean.setWb(cellInfo.getString("word"));
                    dygxxBean.setXsc(cellInfo.getInteger("xsc"));
                    dygxxBean.setYsc(cellInfo.getInteger("ysc"));
                    dygxxBean.setXec(cellInfo.getInteger("xec"));
                    dygxxBean.setYec(cellInfo.getInteger("yec"));
                    dygxxBean.setZbxx(cellInfo.getJSONArray("pos"));
                    dygxxBeans.add(dygxxBean);
                }
                bgxxBean.setDygxx(dygxxBeans);
                bgxxBeans.add(bgxxBean);
            }
        }
        return bgxxBeans;
    }

    /**
     * 封装单页PDF识别后的文本信息
     *
     * @param pdfnrBuilder
     * @param ocrResult
     * @return
     */
    private static List<WbxxBean> packageWbxxInfo(StringBuilder pdfnrBuilder, OcrOutputDTO ocrResult) {
        List<WbxxBean> wbxxBeans = new LinkedList<>();

        JSONArray prismWordsInfo = ocrResult.getPrismWordsInfo();
        if (prismWordsInfo != null && !prismWordsInfo.isEmpty()) {
            for (int i = 0; i < prismWordsInfo.size(); i++) {
                JSONObject prismObj = prismWordsInfo.getJSONObject(i);
                WbxxBean wbxxBean = new WbxxBean();
                wbxxBeans.add(wbxxBean);

                // 文本内容
                wbxxBean.setHwb(prismObj.getString("word"));

                if (pdfnrBuilder != null) {
                    pdfnrBuilder.append(prismObj.getString("word")).append("\n");
                }

                // 文本行号
                wbxxBean.setHxh(prismObj.getInteger("rowId"));
                // 段落号
                wbxxBean.setDlxh(prismObj.getInteger("paragraphId"));
                // 封装行文本坐标信息
                packageHwbzbInfo(prismObj, wbxxBean);
                // 封装字符信息
                packageZfxxInfo(prismObj, wbxxBean);
            }
        }
        return wbxxBeans;
    }

    /**
     * 封装单页PDF识别后的字符信息数据
     *
     * @param prismObj
     * @param wbxxBean
     */
    private static void packageZfxxInfo(JSONObject prismObj, WbxxBean wbxxBean) {
        List<ZfxxBean> zfxxBeans = new LinkedList<>();
        wbxxBean.setZfxx(zfxxBeans);
        JSONArray charInfo = prismObj.getJSONArray("charInfo");
        if (charInfo != null && !charInfo.isEmpty()) {
            for (int i1 = 0; i1 < charInfo.size(); i1++) {
                JSONObject charInfoItem = charInfo.getJSONObject(i1);
                ZfxxBean zfxxBean = new ZfxxBean();
                zfxxBeans.add(zfxxBean);
                zfxxBean.setZ(charInfoItem.getString("word"));
                zfxxBean.setX(charInfoItem.getInteger("x"));
                zfxxBean.setY(charInfoItem.getInteger("y"));
                zfxxBean.setK(charInfoItem.getInteger("w"));
                zfxxBean.setG(charInfoItem.getInteger("h"));
            }
        }
    }

    /**
     * 封装单页PDF识别后的行文本内容
     *
     * @param prismObj
     * @param wbxxBean
     */
    private static void packageHwbzbInfo(JSONObject prismObj, WbxxBean wbxxBean) {
        // 行文本坐标
        List<HwbzbBean> hwbzbBeanList = new LinkedList<>();
        wbxxBean.setHwbzb(hwbzbBeanList);
        JSONArray posArray = prismObj.getJSONArray("pos");
        if (posArray != null && !posArray.isEmpty()) {
            for (int i1 = 0; i1 < posArray.size(); i1++) {
                JSONObject posObj = posArray.getJSONObject(i1);
                hwbzbBeanList.add(JSON.parseObject(posObj.toJSONString(), HwbzbBean.class));
            }
        }
    }
}
